#include "CIatHook.h"

CIatHook g_IatHook;

CIatHook::CIatHook()
{
}

CIatHook::~CIatHook()
{
    RemoveAll();
}

BOOL CIatHook::IatHook(IN const char* szModuleName, IN const char* szFunctionName, IN ULONG_PTR SrcFunc, IN ULONG_PTR NewFunc)
{
    if (!szModuleName || !szFunctionName || !SrcFunc || !NewFunc) return FALSE;

    PIATHOOK lpIatHook;
    HMODULE hModule = GetModuleHandle(NULL);
    PIMAGE_DOS_HEADER lpDosHeader = (PIMAGE_DOS_HEADER)hModule;
    PIMAGE_NT_HEADERS lpNtHeaders = (PIMAGE_NT_HEADERS)((char*)hModule + lpDosHeader->e_lfanew); 
    PIMAGE_DATA_DIRECTORY lpImportEntry = (PIMAGE_DATA_DIRECTORY)
        (lpNtHeaders->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT);
    PIMAGE_IMPORT_DESCRIPTOR lpImportDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)
        ((char*)hModule +  lpImportEntry->VirtualAddress);
    char* szDllName = nullptr;
    ULONG_PTR* pIat = nullptr;
    DWORD dwProtect = 0;
    std::string str;

    str = szFunctionName;
    while (lpImportDescriptor->Name)
    {
        szDllName = (char*)((char*)hModule + lpImportDescriptor->Name);
        if (_stricmp(szDllName, szModuleName) != 0)
        {
            lpImportDescriptor++;
            continue;
        }

        pIat = (ULONG_PTR*)((char*)hModule + lpImportDescriptor->FirstThunk);
        while (*pIat)
        {
            if (*pIat == SrcFunc)
            {
                lpIatHook = new IATHOOK(pIat, SrcFunc);
                if (!lpIatHook) return FALSE;

                VirtualProtect((LPVOID)pIat, IAT_HOOK_ADDRESS_SIZE, PAGE_EXECUTE_READWRITE, &dwProtect);
                *pIat = NewFunc;
                VirtualProtect((LPVOID)pIat, IAT_HOOK_ADDRESS_SIZE, dwProtect, NULL);
                m_IatHookMap[str] = lpIatHook;
                return TRUE;
            }
            pIat++;
        }
        break;
    }

    return FALSE;
}

BOOL CIatHook::Remove(IN const char* szFunctionName)
{
    BOOL bRet = FALSE;
    PIATHOOK lpIatHook;
    std::string str;

    if (szFunctionName)
    {
        str = szFunctionName;
        auto it = m_IatHookMap.find(str);
        if (it != m_IatHookMap.end())
        {
            lpIatHook = it->second;
            *lpIatHook->lpIatAddress = lpIatHook->lpSrcFunc;
            bRet = TRUE;
        }
    }
    return bRet;
}

void CIatHook::RemoveAll()
{
    PIATHOOK lpIatHook;
    for (auto it = m_IatHookMap.begin(); it != m_IatHookMap.end();)
    {
        lpIatHook = it->second;
        *lpIatHook->lpIatAddress = lpIatHook->lpSrcFunc;
        delete lpIatHook;
        it = m_IatHookMap.erase(it);
    }
}
